#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>

#define IMAGE_MAGIC 0x00000803
#define LABEL_MAGIC 0x00000801

typedef struct
{
    uint32_t num_items;
    uint32_t rows;
    uint32_t cols;
    uint8_t *data;
} MNIST_Data;

MNIST_Data read_mnist_images(const char *filename)
{
    MNIST_Data mnist = {0};
    FILE *file = fopen(filename, "rb");

    if (!file)
    {
        perror("Error opening file");
        exit(EXIT_FAILURE);
    }

    /* 读取魔数并转换字节序*/
    uint32_t magic;
    fread(&magic, sizeof(magic), 1, file);
    magic = __builtin_bswap32(magic); /*转换为大端序*/
    if (magic != IMAGE_MAGIC)
    {
        fprintf(stderr, "Invalid MNIST image file format\n");
        exit(EXIT_FAILURE);
    }

    /*读取图像数量、行数、列数*/
    fread(&mnist.num_items, sizeof(mnist.num_items), 1, file);
    fread(&mnist.rows, sizeof(mnist.rows), 1, file);
    fread(&mnist.cols, sizeof(mnist.cols), 1, file);

    /*转换字节序*/
    mnist.num_items = __builtin_bswap32(mnist.num_items);
    mnist.rows = __builtin_bswap32(mnist.rows);
    mnist.cols = __builtin_bswap32(mnist.cols);

    /*计算数据大小并分配内存*/
    size_t data_size = mnist.num_items * mnist.rows * mnist.cols;
    mnist.data = (uint8_t *)malloc(data_size);

    if (!mnist.data)
    {
        perror("Memory allocation failed");
        exit(EXIT_FAILURE);
    }

    /*读取图像数据*/
    fread(mnist.data, sizeof(uint8_t), data_size, file);
    fclose(file);

    return mnist;
}

MNIST_Data read_mnist_labels(const char *filename)
{
    MNIST_Data mnist = {0};
    FILE *file = fopen(filename, "rb");

    if (!file)
    {
        perror("Error opening file");
        exit(EXIT_FAILURE);
    }

    /*读取魔数并转换字节序*/
    uint32_t magic;
    fread(&magic, sizeof(magic), 1, file);
    magic = __builtin_bswap32(magic); /*转换为大端序*/

    if (magic != LABEL_MAGIC)
    {
        fprintf(stderr, "Invalid MNIST label file format\n");
        exit(EXIT_FAILURE);
    }

    /*读取标签数量*/
    fread(&mnist.num_items, sizeof(mnist.num_items), 1, file);
    mnist.num_items = __builtin_bswap32(mnist.num_items);

    /*分配内存*/
    mnist.data = (uint8_t *)malloc(mnist.num_items);

    if (!mnist.data)
    {
        perror("Memory allocation failed");
        exit(EXIT_FAILURE);
    }

    /*读取标签数据*/
    fread(mnist.data, sizeof(uint8_t), mnist.num_items, file);
    fclose(file);

    return mnist;
}

void normalize_images(MNIST_Data *images)
{
    size_t i;
    for (i = 0; i < images->num_items * images->rows * images->cols; i++)
    {
        images->data[i] = (float)images->data[i] / 255.0f;
    }
}

/*数据扁平化函数*/
float *flatten_images(MNIST_Data *images)
{
    size_t total_pixels = images->num_items * images->rows * images->cols;
    float *flattened = (float *)malloc(total_pixels * sizeof(float));

    if (!flattened)
    {
        perror("Memory allocation failed");
        exit(EXIT_FAILURE);
    }
    size_t i;
    for (i = 0; i < total_pixels; i++)
    {
        flattened[i] = (float)images->data[i] / 255.0f;
    }

    return flattened;
}

void save_flattened(const char *filename, float *data, size_t size)
{
    FILE *file = fopen(filename, "wb");
    fwrite(data, sizeof(float), size, file);
    fclose(file);
}

/*主函数*/
int main()
{
    /*读取训练数据*/
    MNIST_Data train_images = read_mnist_images("/home/monica/read_data/train-images.idx3-ubyte");
    MNIST_Data train_labels = read_mnist_labels("/home/monica/read_data/train-labels.idx1-ubyte");

    /* 读取测试数据 */
    MNIST_Data test_images = read_mnist_images("/home/monica/read_data/t10k-images.idx3-ubyte");
    MNIST_Data test_labels = read_mnist_labels("/home/monica/read_data/t10k-labels.idx1-ubyte");

    /*数据归一化*/
    normalize_images(&train_images);
    normalize_images(&test_images);

    /*数据扁平化*/
    float *flat_train = flatten_images(&train_images);
    float *flat_test = flatten_images(&test_images);

    /*flat_train 和 flat_test 可以传递给矩阵运算组*/

    /*保存扁平化后的数据为二进制文件，供后续矩阵运算使用*/
    save_flattened("flat_train.bin", flat_train, train_images.num_items * train_images.rows * train_images.cols);
    save_flattened("flat_test.bin", flat_test, test_images.num_items * test_images.rows * test_images.cols);




    /* 释放内存*/
    free(train_images.data);
    free(train_labels.data);
    free(test_images.data);
    free(test_labels.data);
    free(flat_train);
    free(flat_test);

    return 0;
}